use std::{fs::File, iter::zip, path::Path, sync::Arc, vec};

use embd::Embd;
use embedding_model::EmbeddingModel;
use serde::{Deserialize, Serialize};

use crate::SearchResultItem;

#[derive(Debug, Serialize, Deserialize, PartialEq, PartialOrd)]
pub struct ContentInner {
    pub content: String,
    pub tags: Option<Vec<String>>,
    pub short: Option<String>,
    pub vec: Vec<f32>,
    pub id: usize,
}

#[derive(Debug, Serialize, Deserialize, PartialEq, PartialOrd)]
pub struct ContentInnerUnembeded {
    pub content: String,
    pub tags: Option<Vec<String>>,
    pub short: Option<String>,
    pub id: usize,
}

impl Eq for ContentInner {}

#[derive(Debug, Clone)]
pub struct Content {
    pub inner: Arc<ContentInner>,
}

impl Serialize for Content {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        self.clone_into_inner().serialize(serializer)
    }
}

impl<'de> Deserialize<'de> for Content {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        Ok(Content::from_inner(ContentInner::deserialize(
            deserializer,
        )?))
    }
}
#[derive(Debug, Clone)]
pub struct ContentStoreInner {
    pub contents: Vec<Content>,
}

/// 存储Content
#[derive(Debug, Clone)]
pub struct ContentStore {
    inner: Arc<ContentStoreInner>,
}

/// 临时存储检索到的Content
pub struct ContentCollection {
    pub contents: Vec<ContentCollectionItem>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContentCollectionItem {
    pub content: Content,
    pub search_result: SearchResultItem,
}

impl ContentCollectionItem {
    pub fn from_parts(content: Content, search_result: SearchResultItem) -> Self {
        Self {
            content,
            search_result,
        }
    }
}

impl ContentCollection {
    pub fn len(&self) -> usize {
        self.contents.len()
    }
    pub fn empty() -> Self {
        Self {
            contents: Vec::new(),
        }
    }
}
impl Into<Vec<Content>> for ContentCollection {
    fn into(self) -> Vec<Content> {
        self.contents
            .into_iter()
            .map(|item| item.content)
            .collect::<Vec<_>>()
    }
}

impl IntoIterator for ContentCollection {
    fn into_iter(self) -> Self::IntoIter {
        self.contents.into_iter()
    }

    type Item = ContentCollectionItem;

    type IntoIter = vec::IntoIter<Self::Item>;
}

impl Content {
    pub fn from_inner(innner: ContentInner) -> Self {
        Self {
            inner: Arc::new(innner),
        }
    }
    pub fn content(&self) -> &str {
        self.inner.content.as_str()
    }
    pub fn blank() -> Self {
        Self {
            inner: Arc::new(ContentInner {
                content: "content".into(),
                tags: Some(vec!["tag1".into()]),
                short: Some("short".into()),
                vec: vec![2.1f32, 4.6f32, 5.8f32],
                id: 0,
            }),
        }
    }
    pub fn clone_into_inner(&self) -> ContentInner {
        ContentInner {
            content: self.inner.content.clone(),
            tags: self.inner.tags.clone(),
            short: self.inner.short.clone(),
            vec: self.inner.vec.clone(),
            id: self.inner.id,
        }
    }
}

impl ContentStore {
    pub fn from_inner(inner: ContentStoreInner) -> Self {
        Self {
            inner: Arc::new(inner),
        }
    }
    pub fn from_json_file<P>(path: P) -> anyhow::Result<Self>
    where
        P: AsRef<Path>,
    {
        let rdr = File::open(path)?;
        let contents: Vec<ContentInner> = serde_json::from_reader(rdr)?;
        let contents = contents
            .into_iter()
            .map(|content| Content::from_inner(content))
            .collect::<Vec<_>>();
        Ok(Self::from_inner(ContentStoreInner { contents }))
    }
    pub fn save_json_file<P>(&self, path: P) -> anyhow::Result<()>
    where
        P: AsRef<Path>,
    {
        let mut wtr = File::create(path)?;
        let contents = self
            .inner
            .contents
            .iter()
            .map(|content| content.clone_into_inner())
            .collect::<Vec<_>>();
        serde_json::to_writer(wtr, &contents)?;
        Ok(())
    }
    pub fn blank() -> Self {
        let contents = ContentInner {
            content: "content".into(),
            tags: Some(vec!["tag1".into()]),
            short: Some("short".into()),
            vec: vec![2.1f32, 4.6f32, 5.8f32],
            id: 0,
        };
        let contents = vec![Content::from_inner(contents)];
        Self::from_inner(ContentStoreInner { contents })
    }
}

impl ContentStore {
    pub fn from_unembeded(
        unembeded: Vec<ContentInnerUnembeded>,
        embd: &Embd,
        batch: usize,
    ) -> anyhow::Result<Self> {
        let contents = unembeded
            .chunks(batch)
            .flat_map(|chk| {
                let sentences = chk
                    .iter()
                    .map(|content| content.content.as_str())
                    .collect::<Vec<_>>();
                let embedding = embd.embedding_batch_vec(sentences);
                match embedding {
                    Ok(embedding) => {
                        let res = zip(chk, embedding)
                            .map(|(content, embedding)| {
                                Content::from_inner(ContentInner {
                                    content: content.content.clone(),
                                    tags: content.tags.clone(),
                                    short: content.short.clone(),
                                    vec: embedding,
                                    id: content.id,
                                })
                            })
                            .collect::<Vec<_>>();
                        res
                    }
                    Err(_) => vec![Content::blank()],
                }
            })
            .collect::<Vec<_>>();
        Ok(Self::from_inner(ContentStoreInner { contents }))
    }
}

impl ContentStore {
    pub fn get_by_id(&self, id: usize) -> Option<Content> {
        match self.inner.contents.get(id) {
            Some(s) => Some(s.clone()),
            None => None,
        }
    }
    pub fn get_inner(&self) -> Arc<ContentStoreInner> {
        self.inner.clone()
    }
    pub fn get_collection(&self, ids: &[SearchResultItem]) -> Option<ContentCollection> {
        let collection = ids
            .iter()
            .map(|item| {
                let content = match self.get_by_id(item.id) {
                    Some(s) => s,
                    None => {
                        panic!();
                    }
                };
                ContentCollectionItem::from_parts(content, item.clone())
            })
            .collect::<Vec<_>>();
        let collection = ContentCollection {
            contents: collection,
        };
        Some(collection)
    }
}

impl ContentInnerUnembeded {
    pub fn from_json_file<P>(path: P) -> anyhow::Result<Vec<Self>>
    where
        P: AsRef<Path>,
    {
        let rdr = File::open(path)?;
        let contents: Vec<ContentInnerUnembeded> = serde_json::from_reader(rdr)?;
        Ok(contents)
    }
    pub fn blank() -> Self {
        Self {
            content: "content".into(),
            //tags: Some(vec!["tag1".into()]),
            tags: None,
            short: None,
            //short: Some("short".into()),
            id: 0,
        }
    }
}
